﻿// 
// Exe/Zip - A poorly-documented blob of code to merge zip archives with executables
// 
// http://exezip.codeplex.com/
// 
// Copyright © 2012 Alex Forster. All rights reserved.
// 
// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
// of the Software, and to permit persons to whom the Software is furnished to do
// so, subject to the following conditions:
// 
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.
//

namespace ExeZip
{
	using System;
	using System.Collections.Generic;
	using System.IO;
	using System.Text;
	using Mono.Options;

	using Zip = ICSharpCode.SharpZipLib.Zip;
	using StreamUtils = ICSharpCode.SharpZipLib.Core.StreamUtils;

	public class Program
	{
		private static String Usage = "";

		private static void PrintError( String message )
		{
			Console.ForegroundColor = ConsoleColor.Yellow;
			Console.BackgroundColor = ConsoleColor.Black;

			Console.WriteLine();
			Console.WriteLine( message );

			Console.ResetColor();

			Console.WriteLine( Program.Usage );
		}

		public static void Main( String[] args )
		{
			//args = "-i Original.exe -o Combined.exe ICSharpCode.SharpZipLib.dll.zip".Split( ' ' );

			var opt_verbose = 0;
			var opt_compression = 0;

			String opt_in = null;
			String opt_out = null;

			var resources = new List<FileSystemInfo>();

			var options = new OptionSet
			{
				{ "h|?|help", "Print this message and exit.", val => { if( val != null ) throw new OptionException( "Usage...", "h" ); } },
				{ "v", "Increase the verbosity of diagnostic messages that are emitted while this program runs.", val => { if( val != null ) ++opt_verbose; else --opt_verbose; } },
				{ "i=|in=", "(Required) Specify the input PE file.", val => opt_in = val },
				{ "o=|out=", "(Required) Specify the name of the combined PE/zip output file.", val => opt_out = val },
				{ "compression=", "Specify the output archive's compression level.", (Byte val) => opt_compression = val },
			};

			var usage = new StringBuilder( @"
Exe/Zip: Combine a zip archive with a PE executable, creating both
         a valid PE file and a valid zip file.

Usage:
    ExeZip -h|-?|--help
    ExeZip [-v[v]] (options) (file/folder [file/folder ...])

Examples:
    ExeZip -i In.exe -o Out.exe Logo.png Config.xml ""C:/Sound Effects""
    ExeZip -vv /in=Exe.exe --out ExeAndZip.exe --compression:9 Zip.zip  (*)

(*) Note: if you specify a single zip archive after your options, then that
    archive will be merged with the input executable. Otherwise, the
    files/folders that you specify will be added to a new zip archive,
    and that new archive will be merged with the input executable.

Options:
" );

			options.WriteOptionDescriptions( new StringWriter( usage ) );

			Program.Usage = usage.ToString();

			try
			{
				if( opt_verbose > 1 ) Console.WriteLine( "Parsing options..." );
				
				var unparsed_options = options.Parse( args );

				if( String.IsNullOrEmpty( opt_in ) )
				{
					throw new OptionException( "This option is required.", "i" );
				}

				if( String.IsNullOrEmpty( opt_out ) )
				{
					throw new OptionException( "Output executable is required.", "o" );
				}

				if( opt_verbose > 1 ) Console.WriteLine( "Validating files/folders..." );

				foreach( var resource in unparsed_options )
				{
					if( Directory.Exists( resource ) )
					{
						if( opt_verbose > 1 ) Console.WriteLine( "Folder '{0}' and its contents will be added to the resulting zip archive.", resource );
						resources.Add( new DirectoryInfo( resource ) );
					}
					else if( File.Exists( resource ) )
					{
						if( opt_verbose > 1 ) Console.WriteLine( "File '{0}' will be added to the resulting zip archive.", resource );
						resources.Add( new FileInfo( resource ) );
					}
					else
					{
						if( opt_verbose > 0 ) Console.WriteLine( "Warning: file/folder '{0}' does not exist.", resource );
					}
				}

				if( resources.Count == 0 )
				{
					Program.PrintError( "No files/folders were specified to be included in the resulting zip archive." );
					return;
				}
			}
			catch( OptionException ex )
			{
				Program.PrintError( "Option -" + ex.OptionName + ": " + ex.Message );

				if( args.Length == 0 )
				{
					Console.Write( "No arguments were provided. Press any key to continue..." );
					Console.ReadKey( true );
				}

				return;
			}

			try
			{
				if( opt_verbose > 1 ) Console.WriteLine( "Validating input executable..." );

				opt_in = Path.GetFullPath( opt_in );

				if( File.Exists( opt_in ) == false )
				{
					throw new FileNotFoundException( "File does not exist.", opt_in );
				}

				using( var f = File.OpenRead( opt_in ) )
				{
					// There's no reason I couldn't leave this check out and let you concat a zip
					// archive to the end of, say, a jpeg. But then again, there's no reason you
					// should want to unless you're doing something sneaky ;)

					if( f.ReadByte() != 0x4D || f.ReadByte() != 0x5A )
					{
						if( opt_verbose > 0 ) Console.WriteLine( "Warning: input executable does not begin with 0x5A4D." );
						throw new FileNotFoundException( "File is not a PE executable.", opt_in );
					}
				}
			}
			catch( Exception ex )
			{
				Program.PrintError( "Option -i: " + ex.Message );
				return;
			}

			try
			{
				if( opt_verbose > 1 ) Console.WriteLine( "Validating output executable..." );

				opt_out = Path.GetFullPath( opt_out );
			}
			catch( Exception ex )
			{
				Program.PrintError( "Option -o: " + ex.Message );
				return;
			}

			try
			{
				var input_exe = new FileInfo( opt_in );
				var output_exe = new FileInfo( opt_out );

				if( opt_verbose > 0 ) Console.WriteLine( "Writing a new zip archive to the output executable at position {0}...", input_exe.Length );

				var raw = output_exe.Open( FileMode.OpenOrCreate, FileAccess.ReadWrite, FileShare.ReadWrite );
				raw.SetLength( input_exe.Length );
				raw.Seek( 0, SeekOrigin.End );

				var output_buffer = new Byte[4096];
				var output_stream = new Zip.ZipOutputStream( raw )
				{
					IsStreamOwner = false,
					UseZip64 = Zip.UseZip64.Off
				};

				if( opt_compression > -1 && opt_compression < 10 )
				{
					output_stream.SetLevel( opt_compression );
				}

				if( resources.Count == 1 &&
				    resources[0] is FileInfo &&
				    resources[0].Extension.Equals( ".zip", StringComparison.OrdinalIgnoreCase ) )
				{
					// copy the *contents* of this single zip archive into the output archive

					var original = (FileInfo)resources[0];

					var zip = new Zip.ZipFile( original.FullName );

					for( var i = 0; i < zip.Count; i++ )
					{
						var e = new Zip.ZipEntry( zip[i].Name )
						{
							AESKeySize = zip[i].AESKeySize,
							Comment = zip[i].Comment,
							CompressionMethod = zip[i].CompressionMethod,
							DateTime = zip[i].DateTime,
							DosTime = zip[i].DosTime,
							ExternalFileAttributes = zip[i].ExternalFileAttributes < 0 ? 0 : zip[i].ExternalFileAttributes,
							ExtraData = zip[i].ExtraData,
							Flags = zip[i].Flags < 0 ? 0 : zip[i].Flags,
							HostSystem = zip[i].HostSystem < 0 ? 0 : zip[i].HostSystem,
							IsCrypted = zip[i].IsCrypted,
							IsUnicodeText = zip[i].IsUnicodeText,
							Size = zip[i].Size
						};

						output_stream.PutNextEntry( e );
						StreamUtils.Copy( zip.GetInputStream( zip[i] ), output_stream, output_buffer );
					}
				}
				else
				{
					// copy the specified files into a new zip archive

					foreach( var fsinfo in resources )
					{
						if( fsinfo is DirectoryInfo )
						{
							var resource = fsinfo as DirectoryInfo;

							var dirname = Zip.ZipEntry.CleanName( resource.Name + Path.DirectorySeparatorChar );

							var entry = new Zip.ZipEntry( dirname )
							{
								IsUnicodeText = true,
								DateTime = input_exe.LastWriteTime
							};

							if( opt_verbose > 1 ) Console.WriteLine( "Adding the directory '{0}' to the zip archive...", dirname );

							output_stream.PutNextEntry( entry );
							output_stream.CloseEntry();

							foreach( var dir in resource.GetDirectories( "*", SearchOption.AllDirectories ) )
							{
								dirname = Zip.ZipEntry.CleanName(
									new Uri( resource.FullName )
										.MakeRelativeUri( new Uri( dir.FullName ) )
										.ToString() + Path.DirectorySeparatorChar );

								entry = new Zip.ZipEntry( dirname )
								{
									IsUnicodeText = true,
									DateTime = input_exe.LastWriteTime
								};

								if( opt_verbose > 1 ) Console.WriteLine( "Adding the directory '{0}' to the zip archive...", dirname );

								output_stream.PutNextEntry( entry );
								output_stream.CloseEntry();
							}

							foreach( var file in resource.GetFiles( "*", SearchOption.AllDirectories ) )
							{
								var filename = Zip.ZipEntry.CleanName(
									new Uri( resource.FullName )
										.MakeRelativeUri( new Uri( file.FullName ) )
										.ToString() );

								entry = new Zip.ZipEntry( filename )
								{
									IsUnicodeText = true,
									DateTime = input_exe.LastWriteTime
								};

								if( opt_verbose > 1 ) Console.WriteLine( "Adding the file '{0}' to the ip archive...", filename );

								output_stream.PutNextEntry( entry );

								using( var file_stream = File.OpenRead( file.FullName ) )
								{
									StreamUtils.Copy( file_stream, output_stream, output_buffer );
								}

								output_stream.CloseEntry();
							}
						}
						else if( fsinfo is FileInfo )
						{
							var resource = fsinfo as FileInfo;

							var filename = Zip.ZipEntry.CleanName( resource.Name );

							var entry = new Zip.ZipEntry( filename )
							{
								IsUnicodeText = true,
								DateTime = input_exe.LastWriteTime
							};

							if( opt_verbose > 1 ) Console.WriteLine( "Adding the file '{0}' to the zip archive...", filename );

							output_stream.PutNextEntry( entry );

							using( var file_stream = File.OpenRead( resource.FullName ) )
							{
								StreamUtils.Copy( file_stream, output_stream, output_buffer );
							}

							output_stream.CloseEntry();
						}
					}
				}

				output_stream.Close();

				if( opt_verbose > 0 ) Console.WriteLine( "Writing the input executable to the beginning of the output executable..." );

				raw.Seek( 0, SeekOrigin.Begin );

				using( var file_stream = File.Open( input_exe.FullName, FileMode.Open, FileAccess.Read, FileShare.ReadWrite ) )
				{
					StreamUtils.Copy( file_stream, raw, output_buffer );
				}

				raw.Close();

				output_exe.LastWriteTime = input_exe.LastWriteTime;

				if( opt_verbose > 0 ) Console.WriteLine( "Finished." );
			}
			catch( Exception ex )
			{
				Program.PrintError( ex.Message );
			}
		}
	}
}
